#!/usr/bin/env python3
import datetime

import yaml
import sys
import subprocess
import re
import os
import time
import signal
import getopt


config = {
    'config_paths': ['fan_control.yaml', '/opt/fan_control/fan_control.yaml'],
    'general': {
        'debug': False,
        'interval': 60
    },
    'hosts': []
}
state = {}


class ConfigError(Exception):
    pass


def ipmitool(args, host):
    global state

    cmd = ["ipmitool"]

    cmd += ['-I', 'lanplus']
    cmd += ['-H', host['ip']]
    cmd += ['-U', host['username']]
    cmd += ['-P', host['password']]
    cmd += (args.split(' '))
    if config['general']['debug']:
        print(re.sub(r'-([UP]) (\S+)', r'-\1 ___', ' '.join(cmd)))  # Do not log IPMI credentials
        return 'True'

    try:
        res = subprocess.check_output(cmd, timeout=15)
        return str(res)
    except subprocess.CalledProcessError:
        print("\"{}\" command has returned a non-0 exit code".format(cmd), file=sys.stderr)
        return 'False'
    except subprocess.TimeoutExpired:
        print("\"{}\" command has timed out".format(cmd), file=sys.stderr)
        return 'False'


def set_fan_control(wanted_mode, host):
    global state
    # manual 手动调速：即人工控制或脚本自动控制
    # automatic 自动调速：服务器本身自动控制风速
    if wanted_mode == "manual" or wanted_mode == "automatic":
        if wanted_mode == "manual" and state[host['name']]['fan_control_mode'] == "automatic":
            if not config['general']['debug']:
                print("[{}] Switching to manual mode".format(host['name']))
            # 0x00 代表 禁用 自动调速
            ipmitool("raw 0x30 0x30 0x01 0x00", host)
        elif wanted_mode == "automatic" and state[host['name']]['fan_control_mode'] == "manual":
            if not config['general']['debug']:
                print("[{}] Switching to automatic mode".format(host['name']))
            # 0x01 代表 启用 自动调速
            ipmitool("raw 0x30 0x30 0x01 0x01", host)
            state[host['name']]['fan_speed'] = 0

        state[host['name']]['fan_control_mode'] = wanted_mode


def set_fan_speed(threshold_n, host):
    global state

    wanted_percentage = host['speeds'][threshold_n]
    if wanted_percentage == state[host['name']]['fan_speed']:
        return

    if 5 <= wanted_percentage <= 100:
        wanted_percentage_hex = "{0:#0{1}x}".format(wanted_percentage, 4)
        if state[host['name']]['fan_control_mode'] != "manual":
            set_fan_control("manual", host)
            time.sleep(1)
        if not config['general']['debug']:
            print("[{}] Setting fans speed to {}%".format(host['name'], wanted_percentage))
        # 调整风速：0xff 等于所有风扇，最后一个参数为风扇转速的百分比，用十六进制表示
        # 0x18 你期望风扇运行的转速百分比，原厂的暴力扇应该为12000rpm,这是 24 的16进制
        # 0xff 0x18 = 所有风扇 以 24% 的转速运行，可以更改 0xff 为单独的风扇调速
        ipmitool("raw 0x30 0x30 0x02 0xff {}".format(wanted_percentage_hex), host)
        state[host['name']]['fan_speed'] = wanted_percentage


def parse_config():
    global config
    _debug = config['general']['debug']
    _interval = config['general']['interval']

    config_path = None
    for path in config['config_paths']:
        if os.path.isfile(path):
            config_path = path
    if not config_path:
        raise RuntimeError("Missing or unspecified configuration file.")
    else:
        print("Loading configuration file.")
        _config = None
        try:
            with open(config_path, 'r') as yaml_conf:
                _config = yaml.safe_load(yaml_conf)
        except yaml.YAMLError as err:
            raise err  # TODO: pretty print
        config = _config
        if 'debug' not in list(config['general'].keys()):
            config['general']['debug'] = _debug
        if 'interval' not in list(config['general'].keys()):
            config['general']['interval'] = _interval

        for host in config['hosts']:
            if 'hysteresis' not in list(host.keys()):
                host['hysteresis'] = 0
            if 'name' not in list(host.keys()):
                host['name'] = 'IDRAC-FANS-CONTROLLER'
            if len(host['temperatures']) != 3:
                raise ConfigError('Host "{}" has {} temperature thresholds instead of 3.'.format(host['name'], len(host['temperatures'])))
            if len(host['speeds']) != 3:
                raise ConfigError('Host "{}" has {} fan speeds instead of 3.'.format(host['name'], len(host['speeds'])))
            if 'ip' not in list(host.keys()) or 'username' not in list(host.keys()) or 'password' not in list(host.keys()):
                raise ConfigError('Host "{}" must specify either none or all "host", "username" and "password" values.'.format(host['name']))
            # TODO: check presence/validity of values instead of keys presence only

            if host['name'] in list(state.keys()):
                raise ConfigError('Duplicate "{}" host name found.'.format(host['name']))
            state[host['name']] = {
                'fan_control_mode': 'automatic',
                'fan_speed': 0
            }


def parse_opts():
    global config
    help_str = "fan_control.py [-d] [-c <path_to_config>] [-i <interval>]"

    try:
        opts, _ = getopt.getopt(sys.argv[1:], "hdc:i:", ["help", "debug", "config=", "interval="])
    except getopt.GetoptError as e:
        print("Unrecognized option. Usage:\n{}".format(help_str))
        raise getopt.GetoptError(e)

    for opt, arg in opts:
        if opt in ('-h', '--help'):
            print(help_str)
            raise InterruptedError
        elif opt in ('-d', '--debug'):
            config['general']['debug'] = True
        elif opt in ('-c', '--config'):
            config['config_paths'] = [arg]
        elif opt in ('-i', '--interval'):
            config['general']['interval'] = arg


def graceful_shutdown(signalnum, frame):
    print("Signal {} received, giving up control".format(signalnum))
    for host in config['hosts']:
        set_fan_control("automatic", host)
    sys.exit(0)


def check_hysteresis(temperature, threshold_n, host):
    global state

    # Skip checks if hysteresis is disabled for this host
    if not host['hysteresis']:
        return True

    # Fan speed is higher than it should be or automatic mode is currently enabled
    if (state[host['name']]['fan_speed'] > host['speeds'][threshold_n] or
            state[host['name']]['fan_control_mode'] == 'automatic'):
        # T ≤ (threshold - hysteresis)
        return temperature <= host['temperatures'][threshold_n] - host['hysteresis']

    # Fan speed is lower than it should be, step up immediately and ignore hysteresis
    return True


def compute_fan_speed(temp_average, host):
    global state

    if config['general']['debug']:
        print("[{}] T:{}°C M:{} S:{}%".format(host['name'], temp_average, state[host['name']]['fan_control_mode'], state[host['name']]['fan_speed']))

    # Tavg < Threshold0
    if (
        temp_average <= host['temperatures'][0] and
        check_hysteresis(temp_average, 0, host)
    ):
        set_fan_speed(0, host)

    # Threshold0 < Tavg ≤ Threshold1
    elif (
        host['temperatures'][0] < temp_average <= host['temperatures'][1] and
        check_hysteresis(temp_average, 1, host)
    ):
        set_fan_speed(1, host)

    # Threshold1 < Tavg ≤ Threshold2
    elif (
        host['temperatures'][1] < temp_average <= host['temperatures'][2] and
        check_hysteresis(temp_average, 2, host)
    ):
        set_fan_speed(2, host)

    # Tavg > Threshold2
    elif host['temperatures'][2] < temp_average:
        set_fan_control("automatic", host)


def parse_sensor_data(data):
    # 从一行数据切分为多行数据，得到多个传感器数据
    # 例如：“Temp     | 41.000  | degrees C  | ok    | na     | 3.000  | 8.000  | 95.000    | 100.000   | na ”
    data = data.split(r"\r\n")
    result = []
    # 将每一个传感器数据继续进行切分
    # 并且只保留关心的传感器数据
    sensor_names = ['Intrusion', 'Fan1', 'Fan2', 'Fan3', 'Fan4', 'Fan5', 'Fan6', 'InletTemp', 'ExhaustTemp',
                    'PwrConsumption', 'PowerOptimized', 'Temp']
    for item in data:
        temp = item.replace(' ', '').split('|')
        if temp[0] in sensor_names:
            result.append(temp)  # 去除空格后进行切分
    return result


def main():
    global config
    global state

    print("Starting fan control script.")
    for host in config['hosts']:
        print("[{}] Thresholds of {}°C ({}%), {}°C ({}%) and {}°C ({}%)".format(
                host['name'],
                host['temperatures'][0], host['speeds'][0],
                host['temperatures'][1], host['speeds'][1],
                host['temperatures'][2], host['speeds'][2],
            ))

    while True:
        for host in config['hosts']:
            temps = []

            # ipmitool使用sensor读取CPU温度
            temp = ipmitool("sensor", host)
            if temp == "False":
                continue
            result = parse_sensor_data(temp)
            current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
            print(host['name'] + '\'s sensor data on ' + current_time + ':\n')
            # 将CPU的温度取出来，排除'na'值，传感器名字：Temp
            for item in result:
                print(item)
                if item[0] == 'Temp':
                    try:
                        temps.append(eval(item[1]))
                    except SyntaxError:
                        continue
                    except NameError:
                        continue
            # 计算平均值，进行控制
            temp_average = round(sum(temps)/len(temps))
            compute_fan_speed(temp_average, host)
        # 间隔interval秒再进行监测控制
        time.sleep(config['general']['interval'])


if __name__ == "__main__":
    # Reset fan control to automatic when getting killed
    signal.signal(signal.SIGTERM, graceful_shutdown)

    try:
        try:
            parse_opts()
        except (getopt.GetoptError, InterruptedError):
            sys.exit(1)
        parse_config()
        main()
    finally:
        print("Closing")
